{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff14471a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    DiagramBuilder,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Parser,\n",
    "    MeshcatVisualizer,\n",
    "    StartMeshcat,\n",
    "    MeshcatVisualizerParams,\n",
    "    Simulator,\n",
    "    JointSliders,\n",
    "    InverseKinematics,\n",
    "    RotationMatrix,\n",
    "    Solve,\n",
    "    ContactVisualizerParams,\n",
    "    ContactVisualizer,\n",
    "    GeometrySet,\n",
    "    CollisionFilterDeclaration,\n",
    "    Role,\n",
    "    eq,\n",
    "    RigidTransform,\n",
    ")\n",
    "from manipulation import ConfigureParser, running_as_notebook\n",
    "import pydot\n",
    "from IPython.display import display, Image, SVG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dbf55e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f929f5",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "We will solve inverse kinematics for one arm on a PR2 robot. This robot model has been modified to have convex collision geometries, and to simplify this problem, we have welded (fixed) several joints that are irrelavent to the kinematics of the arms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6afc78f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_env():\n",
    "    \"\"\"Load in models and build the diagram.\"\"\"\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.AddModelsFromUrl(\"package://manipulation/pr2_shelves.dmd.yaml\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph.get_query_output_port(),\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(delete_on_initialization_event=False),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    return diagram, plant, scene_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f40eede",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filterCollsionGeometry(scene_graph, context=None):\n",
    "    \"\"\"Some robot models may appear to have self collisions due to overlapping collision geometries.\n",
    "    This function filters out such problems for our PR2 model.\"\"\"\n",
    "    if context is None:\n",
    "        filter_manager = scene_graph.collision_filter_manager()\n",
    "    else:\n",
    "        filter_manager = scene_graph.collision_filter_manager(context)\n",
    "    inspector = scene_graph.model_inspector()\n",
    "\n",
    "    pr2 = {}\n",
    "    shelves = []\n",
    "    tables = []\n",
    "\n",
    "    for gid in inspector.GetGeometryIds(\n",
    "        GeometrySet(inspector.GetAllGeometryIds()), Role.kProximity\n",
    "    ):\n",
    "        gid_name = inspector.GetName(inspector.GetFrameId(gid))\n",
    "        if \"pr2\" in gid_name:\n",
    "            link_name = gid_name.split(\"::\")[1]\n",
    "            pr2[link_name] = [gid]\n",
    "\n",
    "    def add_exclusion(set1, set2=None):\n",
    "        if set2 is None:\n",
    "            filter_manager.Apply(\n",
    "                CollisionFilterDeclaration().ExcludeWithin(GeometrySet(set1))\n",
    "            )\n",
    "        else:\n",
    "            filter_manager.Apply(\n",
    "                CollisionFilterDeclaration().ExcludeBetween(\n",
    "                    GeometrySet(set1), GeometrySet(set2)\n",
    "                )\n",
    "            )\n",
    "\n",
    "    # Robot-to-self collisions\n",
    "    add_exclusion(\n",
    "        pr2[\"base_link\"],\n",
    "        pr2[\"l_shoulder_pan_link\"]\n",
    "        + pr2[\"r_shoulder_pan_link\"]\n",
    "        + pr2[\"l_upper_arm_link\"]\n",
    "        + pr2[\"r_upper_arm_link\"]\n",
    "        + pr2[\"head_pan_link\"]\n",
    "        + pr2[\"head_tilt_link\"],\n",
    "    )\n",
    "    add_exclusion(\n",
    "        pr2[\"torso_lift_link\"], pr2[\"head_pan_link\"] + pr2[\"head_tilt_link\"]\n",
    "    )\n",
    "    add_exclusion(\n",
    "        pr2[\"l_shoulder_pan_link\"] + pr2[\"torso_lift_link\"],\n",
    "        pr2[\"l_upper_arm_link\"],\n",
    "    )\n",
    "    add_exclusion(\n",
    "        pr2[\"r_shoulder_pan_link\"] + pr2[\"torso_lift_link\"],\n",
    "        pr2[\"r_upper_arm_link\"],\n",
    "    )\n",
    "    add_exclusion(pr2[\"l_forearm_link\"], pr2[\"l_gripper_palm_link\"])\n",
    "    add_exclusion(pr2[\"r_forearm_link\"], pr2[\"r_gripper_palm_link\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba4076c",
   "metadata": {},
   "outputs": [],
   "source": [
    "goal_rotation = RotationMatrix(\n",
    "    [\n",
    "        [1, 0, 0],\n",
    "        [0, -1, 0],\n",
    "        [0, 0, -1],\n",
    "    ]\n",
    ")\n",
    "goal_position = np.array([-0.83, 0.18, 1.4])\n",
    "goal_pose = RigidTransform(goal_rotation, goal_position)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a950a302",
   "metadata": {},
   "source": [
    "### IK for a Mobile Manipulator\n",
    "\n",
    "Given a RigidTransform X_WG, compute a robot configuration placing the left gripper at that pose. We use optimization to solve the IK problem, and we repeatedly solve the program with random initializations until it succeeds. We have implemented a skeleton of the necessary code in the following function, but you must complete several pieces:\n",
    "\n",
    "- Add position and orientation constraints to the gripper frame. The end effector should match the desired pose to within 1mm translation along each axis, and the rotation should be off by no more than 1 degree.\n",
    "- If `fix_base` is True, constrain the base pose $(x,y,\\theta)$ to be equal to `base_pose`.\n",
    "- Add a collision free constraint with [AddMinimumDistanceLowerBoundConstraint](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_inverse_kinematics.html#a2ecd71efd675a7e1a4293adb05c9b9df). The minimum distance between any pair of collision geometries should be at least 1cm.\n",
    "- Compute a random initial guess for the joint angles within the robot's joint limits. You can access the joint limits from the multibody plant, but some of the joints are angle-valued and don't have limits. For these joints, use the range $[-\\pi,\\pi]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "906ce3e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_ik(X_WG, max_tries=10, fix_base=False, base_pose=np.zeros(3)):\n",
    "    diagram, plant, scene_graph = build_env()\n",
    "\n",
    "    gripper_frame = plant.GetFrameByName(\"l_gripper_palm_link\")\n",
    "\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "    sg_context = scene_graph.GetMyContextFromRoot(context)\n",
    "    filterCollsionGeometry(scene_graph, sg_context)\n",
    "\n",
    "    # Note: passing in a plant_context is necessary for collision-free constraints!\n",
    "    ik = InverseKinematics(plant, plant_context)\n",
    "    q_variables = ik.q()  # Get variables for MathematicalProgram\n",
    "    prog = ik.prog()  # Get MathematicalProgram\n",
    "    q_nominal = np.zeros(len(q_variables))\n",
    "    q_nominal[0:3] = base_pose\n",
    "    prog.AddQuadraticErrorCost(\n",
    "        np.eye(len(q_variables)), q_nominal, q_variables\n",
    "    )\n",
    "\n",
    "    # Add your constraints here\n",
    "\n",
    "    for count in range(max_tries):\n",
    "        # Compute a random initial guess here\n",
    "\n",
    "        result = Solve(prog)\n",
    "\n",
    "        if running_as_notebook:\n",
    "            render_context = diagram.CreateDefaultContext()\n",
    "            plant.SetPositions(\n",
    "                plant.GetMyContextFromRoot(render_context),\n",
    "                result.GetSolution(q_variables),\n",
    "            )\n",
    "            diagram.ForcedPublish(context)\n",
    "\n",
    "        if result.is_success():\n",
    "            print(\"Succeeded in %d tries!\" % (count + 1))\n",
    "            return result.GetSolution(q_variables)\n",
    "\n",
    "    print(\"Failed!\")\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "461257ca",
   "metadata": {},
   "source": [
    "First, we show an example where we have fixed the base. It may take many tries for it to solve, or not solve at all! At each iteration, we visualize where the optimizer stopped, so you can see what the failures look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbfbf447",
   "metadata": {},
   "outputs": [],
   "source": [
    "solve_ik(\n",
    "    goal_pose,\n",
    "    max_tries=20,\n",
    "    fix_base=True,\n",
    "    base_pose=np.array([-1.23, 0.05, 0]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69727ae3",
   "metadata": {},
   "source": [
    "When we allow the base to move freely, the inverse kinematics can be solved much more easily!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0436add7",
   "metadata": {},
   "outputs": [],
   "source": [
    "solve_ik(goal_pose, fix_base=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f0ff635",
   "metadata": {},
   "outputs": [],
   "source": [
    "from manipulation.exercises.mobile.test_mobile_base_ik import (\n",
    "    TestMobileBaseIk,\n",
    ")\n",
    "from manipulation.exercises.grader import Grader\n",
    "\n",
    "Grader.grade_output([TestMobileBaseIk], [locals()], \"results.json\")\n",
    "Grader.print_test_results(\"results.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cefcbfb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}